from ctypes import *
#from multiprocessing import Process, Queue
import queue
import time
from threading import Lock,Thread
from fastapi import FastAPI
from fastapi import Request
from fastapi import WebSocket, WebSocketDisconnect
import uvicorn
#from yolo_service import *
import socket
import random
from typing import List
import darknet
import cv2
import time
import io
import struct
import os
import numpy as np
import base64
import json
from jtracer.tracing import init_tracer
import opentracing
import pynng
from PIL import Image
from opentracing.propagation import Format

result_queue = None
darknet_height = None
darknet_width = None
network = None
class_names = None
class_colors = None
init_ready = False
tracer =  opentracing.tracer
def convert2relative(bbox):
    """
    YOLO format use relative coordinates for annotation
    """
    x, y, w, h  = bbox
    return x/darknet_width, y/darknet_height, w/darknet_width, h/darknet_height


def convert2original(image, bbox):
    x, y, w, h = convert2relative(bbox)

    image_h, image_w, __ = image.shape

    orig_x       = int(x * image_w)
    orig_y       = int(y * image_h)
    orig_width   = int(w * image_w)
    orig_height  = int(h * image_h)

    bbox_converted = (orig_x, orig_y, orig_width, orig_height)

    return bbox_converted

class SuperbFrame:
    def __init__(self,dw,dh):
        self.image = None
        self.results = None
        self.darknet_image = darknet.make_image(dw,dh,3)
        self.recv_timestamp = 0
        self.send_timestamp = 0
        self.inference_time = 0
        #self.final_image = None
        self.bytes = None
        self.span = None


def port_is_used(port,ip="0.0.0.0"):
    s = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    try:
        s.connect((ip,port))
        s.shutdown(2)
        return True
    except:
        return False

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        # 存放激活的ws连接对象
        self.active_connections: List[WebSocket] = []

    async def connect(self, ws: WebSocket):
        # 等待连接
        await ws.accept()
        # 存储ws连接对象
        self.active_connections.append(ws)

    def disconnect(self, ws: WebSocket):
        # 关闭时 移除ws对象
        self.active_connections.remove(ws)
    
    async def broadcast(self,message):
        for connection in self.active_connections:
            await connection.send_json(message)

manager = ConnectionManager()

@app.get("/get_port")
def get_port(request:Request):
    print("get_port")
    global network
    global class_colors
    global class_names
    if network is None:
        network,class_names,class_colors = darknet.load_network(
            "./cfg/yolov4-tiny.cfg",
            "./cfg/coco.data",
            "./yolov4-tiny.weights",
            batch_size=1
            )

    global result_queue
    global darknet_width
    global darknet_height
    darknet_width = darknet.network_width(network)
    darknet_height = darknet.network_height(network)
    
    port_tmp = 0
    while True:
        port_tmp = random.randint(int(os.getenv("SUPB_MIN_PORT")),int(os.getenv("SUPB_MAX_PORT")))
        if port_is_used(port_tmp):
            continue
        else:
            break
    result_queue = queue.Queue(maxsize=1)
    input_address = "tcp://0.0.0.0:"+str(port_tmp)
    

    #p0 = Thread(target=keep_recv,args=(input_queue,sock,darknet_width,darknet_height))
    p0 = Thread(target=keep_recv,args=(input_address,darknet_width,darknet_height))
    #p1 = Thread(target=keep_inference,args=(send_queue,input_queue,result_queue))
    #p2 = Thread(target=send_index,args=(send_queue,sock))
    p0.start()
    #p1.start()
    #p2.start()
    global init_ready
    init_ready = True
    return port_tmp # port_tmp is the key for a client

def parse_data(data):
    global tracer
    head_length, msg_length = struct.unpack("ii", data[0:8])
    head_length, msg_length, msg_head, msg = struct.unpack("ii"+ str(head_length) + "s" + str(msg_length) + "s", data)
    
    if head_length > 2: 
        span_dict = json.loads(msg_head)
        span_ctx = tracer.extract(Format.TEXT_MAP, span_dict)
        return span_ctx, msg
    else:
        return None, msg

def keep_recv(input_address,dw,dh):
    sock = pynng.Pair1(recv_timeout=10,send_timeout=10) 
    sock.listen(input_address)
    while True:
        try:
            msg = sock.recv()
        except:
            continue
        #with tracer.start_span('pre-processing') as span:
        newFrame = SuperbFrame(dw,dh)
        newFrame.recv_timestamp = int(time.time()*1000.0)

        # msg handling
        span_ctx, msg_content = parse_data(msg)
        if span_ctx is not None:
            newFrame.span = tracer.start_span('image_procss',child_of=span_ctx)
        
        header = msg_content[0:24]
        hh,ww,cc,tt = struct.unpack('iiid',header)
        newFrame.send_timestamp = int(tt*1000.0)
        #print("pynng time:",newFrame.recv_timestamp-newFrame.send_timestamp)
        hh,ww,cc,tt,ss = struct.unpack('iiid'+str(hh*ww*cc)+'s',msg_content)

        newFrame.image = cv2.cvtColor((np.frombuffer(ss,dtype=np.uint8)).reshape(hh,ww,cc), cv2.COLOR_BGR2RGB)
        darknet.copy_image_from_bytes(newFrame.darknet_image,cv2.resize(newFrame.image,(darknet_width,darknet_height),interpolation=cv2.INTER_LINEAR).tobytes())
        prev_time = time.time()
        newFrame.results = darknet.detect_image(network, class_names, newFrame.darknet_image, thresh=0.2)
        newFrame.inference_time = int((time.time()-prev_time)*1000.0) # s -> ms
        #darknet.free_image(newFrame.darknet_image)
        if newFrame.span is not None:
            index = newFrame.span.get_baggage_item('index')
            newFrame.span.finish()
            try:
                sock.send(index.encode())
                pass
            except:
                pass
        try:
            result_queue.put(newFrame,block=False,timeout=1)
        except:
            continue


@app.websocket("/ws")# user is the received port_tmp
async def stream_handler(wss: WebSocket):
    print("a new websocket connected")
    if init_ready is False:
        return
    try:
        await manager.connect(wss)
        while True:
            try:
                superbFrame = result_queue.get(block=False,timeout=1)
            except:
                continue
            detections_adjusted = []
            for label, confidence, bbox in superbFrame.results:
                bbox_adjusted = convert2original(superbFrame.image, bbox)
                detections_adjusted.append((str(label), confidence, bbox_adjusted))
            image = darknet.draw_boxes(detections_adjusted, superbFrame.image, class_colors)
            cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
            #newFrame.final_image = image
            img = Image.fromarray(image).resize((320,240))# final size sent to websocket
            img_byte_arr = io.BytesIO()
            img.save(img_byte_arr, format='PNG')
            img_byte_arr.seek(0)
            bytes = base64.b64encode(img_byte_arr.read()).decode()
            send1_time = int(time.time()*1000.0)
            payload = {"img": "data:image/png;base64,%s"%(bytes),"send0_time":superbFrame.send_timestamp,"recv_time":superbFrame.recv_timestamp,"send1_time":send1_time}
            await manager.broadcast(payload)
    except WebSocketDisconnect:
        if wss is not None:
            await manager.disconnect(wss)
    except Exception as e:
        if wss is not None:
            await manager.disconnect(wss)
        print(e)
    return
            
if __name__ == "__main__":
    tracer = init_tracer("image-process")
    uvicorn.run("darknet_websocket_demo:app",host="0.0.0.0",port=int(os.getenv("SUPB_SERVICE_PORT")),log_level="info")
    